package langnstats.project.languagemodel.interpolation;

import java.io.IOException;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import langnstats.project.Global;
import langnstats.project.languagemodel.srilm.NGramLanguageModel;
import langnstats.project.languagemodel.srilm.SriLMParameters.DiscountingMethod;
import langnstats.project.lib.LanguageModel;
import langnstats.project.lib.WordType;
import langnstats.project.lib.crossvalidation.CrossValidationTokenSet;
import langnstats.project.lib.crossvalidation.TrainTokens;

public class LinearInterpolation implements LanguageModel {
	/**
	 * 
	 */
	private static final long serialVersionUID = -6577275223492289950L;
	private transient double[][][] model_probs;
	private List<LanguageModel> lmodels;
	private double[] lambda;
	private double[] fractions;
	private double stopLL;

	
	private double sum_logprobs;
	private double[] prob_components;
	
	public String getDescription() {
		String desc = "Interpolation among \n";
		
		for (int i=0; i<lmodels.size(); i++) {
			desc += "\t" + lmodels.get(i).getDescription();
		}
		return desc;
	}
	
	public void prepare (WordType[] wordType) {
		return;
	}
	
	public double[] predict(WordType wordType) {
		double[] predicts;
		double[] ret = new double[WordType.vocabularySize()];
		int i, j;
		
		for (j=0; j<ret.length; j++) {
			ret[j] = 0;
		}
		
		for (i=0; i<lmodels.size(); i++) {
			predicts = lmodels.get(i).predict(wordType);
			for (j=0; j<ret.length; j++) {
				ret[j] += lambda[i] * predicts[j]; 
			}
		}
		
		return ret;
	}
	
	@Override
	public LanguageModel clone() {
		return this;
	}
	
	public void train(TrainTokens testTokens) {

		lmodels=Global.getAllModels();
		int i, j;
		int nmodels = lmodels.size();
		boolean first = true;
		LanguageModel model;
		Iterator<LanguageModel> it = lmodels.iterator();
		while(it.hasNext()){
			it.next().train(testTokens);
		}
		double ll = 0.0;
		System.err.println("Training Language Model For interpolation:" + getDescription());
		WordType[] tokenArray = testTokens.getTokenArray();
		model_probs = new double[nmodels][][];
		
		WordType wordType;
		Global.debugOut.println("Start predicting! : " + ll);
		for (i=0; i<nmodels; i++) {
			System.err.println(i);
			model = lmodels.get(i);
			model_probs[i] = new double[tokenArray.length][];
			wordType = null;
			
			for (j=0; j<tokenArray.length ; j++) {
				// read model to model_probs
				model_probs[i][j] = model.predict(wordType);
				wordType = tokenArray[j];
				
			}
		}
		
		lambda = new double[nmodels];
		fractions = new double[nmodels];
		prob_components = new double[nmodels];
		                            
		
		// initialize lambda
		for (i=0; i<nmodels; i++) {
			lambda[i] = 1.0 / nmodels;
		}
		Global.debugOut.println("Start tuning! : " + ll);
		// training
		while ((first == true) || (Math.abs(stopLL - ll) > 0.00001)) {
			stopLL = ll;
			if (first == false) {
				double total_nonfixed_lambdas = 0.0;
				double total_nonfixed_fractions = 0.0;
				// assign new lambda
				for (i=0; i<nmodels; i++) {
					total_nonfixed_lambdas += lambda[i];
					total_nonfixed_fractions += fractions[i];
				}
				
				for (i=0; i<nmodels; i++) {
					lambda[i] = (fractions[i] / tokenArray.length);
				}
			}
			
			ll = eval(nmodels, tokenArray);
			first = false;
			Global.debugOut.println("Log-Likelihood : " + ll + " Stop LL :" + stopLL);
			
		}
		Global.debugOut.println("Log-Likelihood : " + ll);
		printLambda(nmodels);
		System.err.println("Training Language Model OK!" + getDescription());
	}
	
	private void printLambda(int nmodels) {
		for (int i=0; i<nmodels; i++) {
			Global.debugOut.println("lambda[" + i + "] = " + lambda[i]);			
		}
	}
	
	private double eval(int nmodels, WordType[] tokenArray) {
		int i, j;
		double total_prob;
		
		sum_logprobs = 0.0;
		for (i=0; i<nmodels; i++) {
			fractions[i] = 0.0;
		}
		
		total_prob = 0.0;
		for (j=0; j<tokenArray.length; j++) {
			for (i=0; i<nmodels; i++) {
				prob_components[i] = lambda[i] * model_probs[i][j][tokenArray[j].getIndex()];
				total_prob += prob_components[i];
			}
			
			for (i=0; i<nmodels; i++) {
				fractions[i] += prob_components[i] / total_prob;
			}
			sum_logprobs += Global.log(2, total_prob);
			total_prob = 0.0;
		}
		
		//Global.debugOut.println("Stop LL = " + stopLL + "\tLog-likelihood = " + sum_logprobs/tokenArray.length);
		return sum_logprobs/tokenArray.length;
	}
	
	public double test(TrainTokens testTokens) {
		int i, j;
		int nmodels;
		LanguageModel model;
		WordType[] tokenArray;
		WordType wordType;
		double[] predicts;
		double sum_logprobs;
		double min_logprob = -10000;
		
		nmodels = lmodels.size(); 
		tokenArray = testTokens.getTokenArray();
		
		for (i=0; i<nmodels; i++) {
			wordType = null;
			model = lmodels.get(i);
			sum_logprobs = 0;
			for (j=0; j<tokenArray.length; j++) {
				predicts = model.predict(wordType);
				wordType = tokenArray[j];
				sum_logprobs += Global.log(2, predicts[wordType.getIndex()]);
			}
			Global.debugOut.println(model.getDescription() + ": " + sum_logprobs/tokenArray.length);
			min_logprob = Math.max(min_logprob, sum_logprobs/tokenArray.length);
		}
		
		return min_logprob;
	}
	
	
	/*private static NGramLanguageModel[] buildNgram(){
		NGramLanguageModel[] res = new NGramLanguageModel[2];
		int[] discount = new int[3];
		int[] cutoff = new int[3];
		cutoff[0] = 0;
		cutoff[1] = 1;
		cutoff[2] = 2;
		discount[0] = discount[1] = discount[2] = 7;
		res[0] = new NGramLanguageModel(DiscountingMethod.Keyser_Ney,3,cutoff,discount);
		discount = new int[2];
		cutoff = new int[2];
		cutoff[0] = 0;
		cutoff[1] = 1;
		discount[0] = discount[1]  = 7;
		res[1] = new NGramLanguageModel(DiscountingMethod.Good_Turing,2,cutoff,discount);		
		return res;
		
	}*/
	
	private static NGramLanguageModel[] buidlNgramGood(){
		NGramLanguageModel[] res = new NGramLanguageModel[8];
		int[] discount = new int[3];
		int[] cutoff = new int[3];
		cutoff[0] = 0;
		cutoff[1] = 3;
		cutoff[2] = 3;
		discount[0] = discount[1] = discount[2] = 7;
		res[0] = new NGramLanguageModel(DiscountingMethod.Good_Turing,3,cutoff,discount);
		
		cutoff = new int[3];
		cutoff[0] = 0;
		cutoff[1] = 2;
		cutoff[2] = 3;
		res[1] = new NGramLanguageModel(DiscountingMethod.Good_Turing,3,cutoff,discount);

		cutoff = new int[3];
		cutoff[0] = 0;
		cutoff[1] = 2;
		cutoff[2] = 2;
		res[2] = new NGramLanguageModel(DiscountingMethod.Good_Turing,3,cutoff,discount);

		cutoff = new int[3];
		cutoff[0] = 0;
		cutoff[1] = 4;
		cutoff[2] = 3;
		res[3] = new NGramLanguageModel(DiscountingMethod.Good_Turing,3,cutoff,discount);
		discount = new int[4];
		cutoff = new int[4];
		cutoff[0] = 0;
		cutoff[1] = 2;
		cutoff[2] = 2;
		cutoff[3] = 2;
		discount[0] = discount[1] = discount[2] = discount[3] = 7;
		res[4] = new NGramLanguageModel(DiscountingMethod.Good_Turing,4,cutoff,discount);
		cutoff = new int[4];
		cutoff[0] = 0;
		cutoff[1] = 4;
		cutoff[2] = 3;
		cutoff[3] = 2;
		discount[0] = discount[1] = discount[2] = discount[3] = 7;
		res[5] = new NGramLanguageModel(DiscountingMethod.Good_Turing,4,cutoff,discount);
		cutoff = new int[4];
		cutoff[0] = 0;
		cutoff[1] = 3;
		cutoff[2] = 3;
		cutoff[3] = 3;
		discount[0] = discount[1] = discount[2] = discount[3] = 7;
		res[6] = new NGramLanguageModel(DiscountingMethod.Good_Turing,4,cutoff,discount);
		cutoff = new int[4];
		cutoff[0] = 0;
		cutoff[1] = 5;
		cutoff[2] = 4;
		cutoff[3] = 3;
		discount[0] = discount[1] = discount[2] = discount[3] = 7;
		res[7] = new NGramLanguageModel(DiscountingMethod.Good_Turing,4,cutoff,discount);

		return res;
		
	}

	/**
	 * @param args
	 */
	public static void main(String[] args) throws IOException {
		System.out.println("Starting LinearInterpolation::main()");
		
		DateFormat dateFormat = new SimpleDateFormat("yyyyMMdd_HHmmss");
		long startTime = System.currentTimeMillis();
		Global.debugOut.println("Starting JLMTools...");
		
		Global.debugOut.println("Preparing training-evaluation set...");
		WordType[] trainATokens = WordType.parse(Global.getTrainA());
		WordType[][] tokens = CrossValidationTokenSet.divide(trainATokens, 10);
//		CrossValidationTokenSet cvts = new CrossValidationTokenSet(trainATokens, 10);
		
		WordType[] trainTokens = tokens[0];
		WordType[] evalTokens = tokens[1];
				
		Global.debugOut.println("Creating Language Models...");
		
		LinearInterpolation li = new LinearInterpolation();
		//NGramLanguageModel[] models = buidlNgram();
		List<LanguageModel> models = Global.getAllModels();
		
		int modelIndex = models.size();
		
		Global.debugOut.println("Training Language Models... ("+modelIndex+")");
//		CrossValidationResult[] resultArray = cvts.crossValidate(models);
//		Arrays.sort(resultArray);
//		Global.debugOut.println(resultArray[resultArray.length-1].makeReport());
		
//		TrainTokens tt = cvts.getTrainTokens(0);
		
		for (int i = 0; i<modelIndex; i++) {
			LanguageModel model = models.get(i);
			model.train(new TrainTokens(trainTokens));
//			models.get(i).train(tt); 
			Global.debugOut.println(dateFormat.format(System.currentTimeMillis()) + ": Train Finished -- " + models.get(i).getDescription());
			li.addLanguageModel(model);
		}
		Global.debugOut.println("Testing Langugage Models...");
		
		double min_ll = li.test(new TrainTokens(evalTokens));
		li.setStopLL(min_ll);

		Global.debugOut.println(dateFormat.format(System.currentTimeMillis()) + ": Train Linear Interpolation");
		li.train(new TrainTokens(evalTokens));
		
		Global.debugOut.println("Program ended normally!");
		Global.debugOut.println("It took " + dateFormat.format(System.currentTimeMillis()-startTime));
	}

	public double getStopLL() {
		return stopLL;
	}
	
	public void setStopLL(double ll) {
		this.stopLL = ll;
	}

	public LinearInterpolation() {
		super();
		lmodels = new ArrayList<LanguageModel>();
	}

	public void addLanguageModel(LanguageModel model) {
		lmodels.add(model);
	}

}
